from PIL import ImageGrab, Image, ImageFilter
import os
import time
import win32com.client as win32

def capture_screen(sheet, start_cell, end_cell):
    """
    捕获Excel表中指定范围的屏幕截图。
    参数:
        sheet: 包含要捕获范围的Excel表sheet对象。
        start_cell (str): 范围的起始单元格（例如，'A1'）。
        end_cell (str): 范围的结束单元格（例如，'B2'）。
    返回:
        img: 捕获范围的图像对象。
    注意:
        该函数将指定范围复制为图片并从剪贴板中检索。
        由于剪贴板操作可能是异步的，因此该函数会等待直到图像可用。
    """
    selected_range = sheet.Range(f"{start_cell}:{end_cell}")
    selected_range.CopyPicture(Appearance=1, Format=2)
    img = ImageGrab.grabclipboard()
    # 获取剪切板内容为异步操作，所以需要等待一段时间
    while img is None:
        time.sleep(0.1)
    
    return img


def get_workbook(file_path, debug=False):
    """
    打开并返回指定路径的Excel工作簿对象。
    参数:
    file_path (str): Excel文件的路径。
    debug (bool): 是否以调试模式启动Excel应用程序。默认为False。
    返回:
    wb: 如果成功打开工作簿，返回工作簿对象；否则返回None。
    异常:
    Exception: 如果打开工作簿失败，将捕获异常并打印错误信息。
    """
    
    # 启动Excel应用程序
    excel = win32.Dispatch("Excel.Application")
    # 设置是否可见, debug为True时可见
    excel.Visible = debug
    # 尝试打开工作簿
    try:
        wb = excel.Workbooks.Open(os.path.abspath(file_path))
    except Exception as e:
        print(f"打开工作簿失败: {e}")
        return None
    
    return wb


def format_image(img, width):
    ori_width, ori_height = img.size

    # 计算缩放比例
    ratio = width / ori_width

    # 缩放图像
    new_size = (width, int(ori_height * ratio))

    return img.resize(new_size)


def splice_images(images, direction="vertical"):
    """
    拼接图像。
    参数:
        images (list): 要拼接的图像列表。
        direction (str): 拼接方向。可选值为"horizontal"（水平）或"vertical"（垂直）。默认为"horizontal".
    返回:
        img: 拼接后的图像。
    注意:
        该函数将按照指定方向拼接图像，使用该函数前需确保所有图像在拼接方向的另一维度上尺寸相同。
        使用本函数拼接的图像会在各图像间留有20像素的白色间隔。
    """
    # 获取图像尺寸
    if direction == "horizontal":
        width = sum([img.width for img in images]) + 20 * (len(images) - 1)
        height = max([img.height for img in images])
    elif direction == "vertical":
        width = max([img.width for img in images])
        height = sum([img.height for img in images]) + 20 * (len(images) - 1)
    else:
        raise ValueError("direction参数应为'horizontal'或'vertical'。")
        return None
    
    # 创建新图像
    img = Image.new("RGB", (width, height), "white")
    # 拼接图像
    if direction == "vertical":
        y_offset = 0
        for i, image in enumerate(images):
            img.paste(image, (0, y_offset))
            y_offset += image.height + 20
    elif direction == "horizontal":
        x_offset = 0
        for i, image in enumerate(images):
            img.paste(image, (x_offset, 0))
            x_offset += image.width + 20
    
    return img
    

def save_image(img, file_path):
    """
    保存图像到指定路径。
    参数:
        img: 要保存的图像对象。
        file_path (str): 保存路径。
    """
    img =img.filter(ImageFilter.SHARPEN)
    img.save(file_path, "PNG")


def main():
    # 打开工作簿
    wb = get_workbook("sample.xlsx", True)
    if wb is None:
        print("打开工作簿失败。")
        return
    
    # 获取info工作表
    info_sheet = wb.Worksheets("info")
    if info_sheet is None:
        print("未找到info工作表，请检查工作簿。")
        return
    
    # info sheet定义了该wb的基本信息，
    # 其中type可选值为：dr(日报)，bt(节拍)，an(异常)，vs(车型统计)，pl(计划)
    # start和end表示该sheet的范围
    # father_plan仅对type为dr类型有效，表示该日报所属的计划
    # date_row，step_col仅对type为pl类型有效，表示日期所在行和工序所在列
    # info_cols = ["name", "type", "start", "end", "father_plan", "date_row", "step_col"]
    col_idx = {
        "name":None,
        "type":None,
        "start":None,
        "end":None,
        "father_plan":None,
        "date_row":None,
        "step_col":None
    }

    # 获取最大不为空的列数
    max_col = info_sheet.UsedRange.Columns.Count
    for col in range(1, max_col + 1):
        if info_sheet.Cells(1, col).Value in col_idx.keys():
            col_idx[info_sheet.Cells(1, col).Value] = col

    # 获取所有类型为dr的行号
    dr_rows = []
    for row in range(2, info_sheet.UsedRange.Rows.Count + 1):
        if info_sheet.Cells(row, col_idx["type"]).Value == "dr":
            dr_rows.append(row)
    
    # 获取所有类型为dr的sheet，及其start和end
    dr_sheets = []
    for row in dr_rows:
        dr_sheets.append((info_sheet.Cells(row, col_idx["name"]).Value, info_sheet.Cells(row, col_idx["start"]).Value, info_sheet.Cells(row, col_idx["end"]).Value))
    
    # 获取所有类型为dr的sheet的截图
    dr_images = []
    for name, start, end in dr_sheets:
        img = capture_screen(wb.Worksheets(name), start, end)
        dr_images.append(format_image(img, 1090))
    
    # 拼接所有日报图像
    dr_img = splice_images(dr_images)

    # 保存日报图像
    save_image(dr_img, "dr.png")

    # 关闭工作簿
    wb.Close(True)


if __name__ == "__main__":
    main()